package com.hsh.db;

import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.db.DBInputFormat;

public class MySqlInputFormat extends InputFormat<LongWritable, MySqlDBWriteAble> {
	
	Log LOG = LogFactory.getLog(MySqlInputFormat.class);

	@Override
	public List<InputSplit> getSplits(JobContext context) throws IOException,
	InterruptedException {
		Configuration conf = context.getConfiguration();
		String tableName = conf.get("tableName");
		int totalCount = getTotal(tableName);
		int mapCount = conf.getInt("mapred.map.tasks",1);
		
		int splitCount = 1;
		int splitSize = totalCount;
		if(totalCount > mapCount){
			splitCount = totalCount % mapCount == 0 ? totalCount/mapCount : totalCount/mapCount + 1;
			splitSize = totalCount/mapCount;
		}
		List<InputSplit> inputSplits = new ArrayList<InputSplit>(splitCount);
		for(int i=0; i<splitCount; ++i){
			DBInputFormat.DBInputSplit inputSplit = new DBInputFormat.DBInputSplit(i * splitCount, i * splitCount + splitSize);
			inputSplits.add(inputSplit);
		}
		
		return inputSplits;
	} 

	@Override
	public RecordReader<LongWritable, MySqlDBWriteAble> createRecordReader(
			InputSplit split, TaskAttemptContext context) throws IOException,
			InterruptedException {
		return new MyRecordReader(split);
	}

	private int getTotal(String tableName){
		String sql = "SELECT COUNT(0) FROM " + tableName ;
		
		Connection conn = null;
		ResultSet rs = null;
		try {
			Class.forName("com.mysql.jdbc.Driver");
			conn = DriverManager.getConnection("jdbc:mysql://localhost:3306/test", "root", "root");
			Statement stmt = conn.createStatement();
			rs = stmt.executeQuery(sql);
			return rs.getInt(0);
		} catch (Exception e) {
			if(rs != null)
				try {
					rs.close();
					if(conn != null)
						conn.close();
				} catch (SQLException e1) {
					e1.printStackTrace();
				}
		}
		return 0;
	}
}
